from pathlib import Path
from typing import Iterable, Any
import re
from pdfminer.high_level import extract_pages

import PyPDF2
from reportlab.lib.pagesizes import letter
from reportlab.lib.units import inch
from reportlab.pdfgen import canvas
from reportlab.pdfbase.ttfonts import TTFont
from reportlab.pdfbase import pdfmetrics
from datetime import datetime
import os 

from typing import List
from pikepdf import Pdf,Page,Rectangle


'''
向目标pdf文件批量添加水印
target_pdf_path:目标pdf文件路径+文件名
watermark_pad_path:水印pdf文件路径+文件名
nrow:水印平铺的行数
ncol:水印平铺的列数
skip_pages:需要跳过不添加水印的页数

'''
class Watermark:
    def __init__(self, 
                skip_pages:List[int] = [],
            ) -> None:        
        """
        为PDF文件添加水印。

        参数:
        watermark_pdf_path (str): 水印PDF文件的路径。
        nrow (int): 每行添加的水印数量。
        ncol (int): 每列添加的水印数量。
        skip_pages (List[int], 可选): 需要跳过不添加水印的页面编号列表，默认为空列表。

        返回:
        None
        """
        self.watermark_pdf_path = '/root/k8s-vue/exam-platform-of-big-data/工具函数/watermark.pdf'
        self.watermark_score_pdf_path = '/root/k8s-vue/exam-platform-of-big-data/工具函数/watermark_score.pdf'
        self.nrow = 1
        self.ncol = 1
        self.skip_pages = skip_pages

    def  generator_score_postions_watermark(self,  postions, page, page_width, page_height) -> None:
        # Register Chinese font
        pdfmetrics.registerFont(TTFont('song', 'simsun.ttc'))

        c = canvas.Canvas('/root/k8s-vue/exam-platform-of-big-data/工具函数/watermark_score.pdf', pagesize=(page_width,page_height))


        for pos in postions:
            if pos['page'] == page:
                #计算水印在目标页面上的位置并添加
                # 颜色改成红色
                c.setFillColorRGB(1, 0, 0)
                # 字体加粗
                c.setFont("song", 10, 2)  # Use the registered Chinese font
        
                x1,y1,x2,y2 = pos['pos'].strip().split(' ')
                c.setFont("song", 10)  # Use the registered Chinese font
                score = pos['score']
                max_score = pos['max_score']
                c.drawString(float(x1),float(y2)-10, f'{score}/{max_score}')

        c.save()
        return c
        
        
    def  generator_score_watermark(self, teacher, score, center_x=100, center_y=100) -> None:
        # Register Chinese font
        pdfmetrics.registerFont(TTFont('song', 'simsun.ttc'))

        c = canvas.Canvas(self.watermark_pdf_path, pagesize=(200,200))
        
        # Draw watermark circle
        c.setLineWidth(width=4)
        c.setStrokeColorRGB(0.1, 0.6, 0.8)
        c.setFillColorRGB(0.1, 0.6, 0.8)
        c.circle(x_cen=center_x, y_cen=center_y, r=60)
        
        c.setLineWidth(width=1)
        c.setStrokeColorRGB(0.1, 0.6, 0.8)
        c.setFillColorRGB(0.1, 0.6, 0.8)
        c.circle(x_cen=center_x, y_cen=center_y, r=55)

        # Add watermark text in Chinese
        c.setFont("song", 10)  # Use the registered Chinese font
        c.drawString(center_x-20, center_y+10, '成绩：{}'.format(str(score)))
        c.drawString(center_x-30, center_y-10, '阅卷人：{}'.format(teacher))
        c.drawString(center_x-20, center_y+30, '★ ★ ★')

        # Add timestamp
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        c.setFont("song", 9)
        c.drawString(center_x-40, center_y-30, f"{timestamp}")

        c.save()
        return c
        
    def add_score_watermark(self, pdf_path ) -> None:

        
        #打开目标PDF文件
        target_pdf = Pdf.open(pdf_path, allow_overwriting_input=True)

        #打开水印PDF文件并提取第一页作为水印
        watermark_pdf = Pdf.open(self.watermark_pdf_path)
        watermark_page = watermark_pdf.pages[0]

        #遍历目标PDF的每一页，根据行列数批量添加水印
        for idx,target_page in enumerate(target_pdf.pages):
            if idx in self.skip_pages:  # 如果当前页码在skip_pages列表中，则跳过
                continue
            for x in range(self.ncol):
                for y in range(self.nrow):
                    #计算水印在目标页面上的位置并添加
                    target_page.add_overlay(
                        watermark_page,
                        Rectangle(
                            # 右上角水印左下角X坐标，通常设为距离页面右侧边缘的小距离（如10单位）
                            target_page.trimbox[2] - watermark_page.trimbox[2] ,
                            # 右上角水印左下角Y坐标，通常设为距离页面顶部边缘的小距离（如10单位）
                            target_page.trimbox[3] - watermark_page.trimbox[3] ,
                            # 右上角水印右上角X坐标，等于左下角X坐标加上水印宽度
                            target_page.trimbox[2] ,
                            # 右上角水印右上角Y坐标，等于左下角Y坐标加上水印高度
                            target_page.trimbox[3] 
                        )
                    )
        #保存添加了水印的PDF文件，并在原文件名基础上添加"已添加水印"的后缀
        target_pdf.save(pdf_path)

    def add_score_watermark_by_postions(self, pdf_path, postions ) -> None:

        
        #打开目标PDF文件
        target_pdf = Pdf.open(pdf_path, allow_overwriting_input=True)

        #打开水印PDF文件并提取第一页作为水印
        watermark_pdf = Pdf.open(self.watermark_score_pdf_path)
        watermark_page = watermark_pdf.pages[0]

        #遍历目标PDF的每一页，根据行列数批量添加水印
        for idx,target_page in enumerate(target_pdf.pages):
            flag = False
            print('当前页：',idx)
            box = target_page.trimbox
            if idx in self.skip_pages:  # 如果当前页码在skip_pages列表中，则跳过
                continue
            for pos in postions:
                if pos['page'] == idx+1:
                    flag = True
                    #计算水印在目标页面上的位置并添加
                    #计算水印在目标页面上的位置并添加
                    index = pos['pos'].strip().split(' ')
                    index = [p for p in index if p!='']
                    x1,y1,x2,y2 = index
                    x1 = target_page.trimbox[2]-35 
                    # y1 = int(y1) - 10
                    x2 = target_page.trimbox[2]-10
                    # y2 = int(y2) + 10
                    pos['pos'] = f'{x1} {y1} {x2} {y2}'
            if flag:
                self.generator_score_postions_watermark(postions, idx+1, box[2], box[3])
                
                #打开水印PDF文件并提取第一页作为水印
                watermark_pdf = Pdf.open(self.watermark_score_pdf_path)
                watermark_page = watermark_pdf.pages[0]
                
                target_page.add_overlay(
                    watermark_page,
                    Rectangle(0,0,box[2],box[3])
                )
                # os.remove(self.watermark_score_pdf_path)

        #保存添加了水印的PDF文件，并在原文件名基础上添加"已添加水印"的后缀
        target_pdf.save(pdf_path.replace('.pdf','_with_watermark.pdf'))


    def make_score_watermark(self, pdf_path, teacher, score ):
        self.generator_score_watermark(teacher, score)
        self.add_score_watermark(pdf_path)
        os.remove(self.watermark_pdf_path)
        
    
    def make_score_watermark_by_postions(self, pdf_path, teacher, score, postions ):
        self.generator_score_watermark(teacher,score)
        self.add_score_watermark_by_postions(pdf_path,  postions)
        if os.path.exists(self.watermark_pdf_path):
            os.remove(self.watermark_pdf_path)
        



class PdfContentPosition:
    #  x1  y1  x2  y2, 坐标原点在左下角。(x1, y1)为左下角坐标，(x2, y2)为右上角坐标。
    def __init__(self):

        pass
    def get_content_pos(self, pdf_path):
        self.pdf_path = pdf_path
        self.cur_page = 0
        self.word_pos = []

        path = Path(pdf_path).expanduser()
        pages = extract_pages(path)
        
        self._show_ltitem_hierarchy(pages)
        return self.word_pos
    def _match_code(self, text):
        
        pattern = r'(代码_[0-9.]*)[\(（ ]*(\d+)*[分)）:：]*'
        return re.findall(pattern, text)
    
    def _match_question_scores(self,text):
        pattern = r'第(\S)题[\. ]*[\(（](\d+)\s*分\)'
        # matches = re.findall(pattern, text)
        return re.findall(pattern, text)

    def _show_ltitem_hierarchy(self, o: Any, depth=0):
        """Show location and text of LTItem and all its descendants"""
        # 如果要查看输出的内容，取消注释下面的3个print
        if depth == 0:
            # print('element                        x1  y1  x2  y2   text')
            # print('------------------------------ --- --- --- ---- -----')
            pass
        if depth == 1:
            self.cur_page+=1
        text = self._get_optional_text(o)
        pos = self._get_optional_bbox(o)
        
        # print(
        #     f'{self._get_indented_name(o, depth):<30.30s} '
        #     f'{pos} '
        #     f'{text}'
        # )
        if depth == 2:
            title_match = self._match_question_scores(text)
            if title_match:
                self.word_pos.append({'page':self.cur_page, 'pos':pos, 'text':text, 'type':'title','score':'', 'max_score':title_match[0][1]})
            else :
                code_match = self._match_code(text)
                if code_match:
                    self.word_pos.append({'page':self.cur_page, 'pos':pos, 'text':text, 'type':'code', 'score':'', 'max_score':code_match[0][1]})
            return 
        if isinstance(o, Iterable):
            for i in o:
                self._show_ltitem_hierarchy(i, depth=depth + 1)

    def _get_indented_name(self, o: Any, depth: int) -> str:
        """Indented name of LTItem"""
        return '  ' * depth + o.__class__.__name__

    def _get_optional_bbox(self, o: Any) -> str:
        """Bounding box of LTItem if available, otherwise empty string"""
        if hasattr(o, 'bbox'):
            return ''.join(f'{i:<4.0f}' for i in o.bbox)
        return ''

    def _get_optional_text(self, o: Any) -> str:
        """Text of LTItem if available, otherwise empty string"""
        if hasattr(o, 'get_text'):
            return o.get_text().strip()
        return ''
if __name__ == '__main__':
    # 使用示例
    pdf_path = '/nfs/bigdata2024/2/homework/2/stuSubmit/20204249/test3.pdf'
    pdfPosGetter = PdfContentPosition()
    pos = pdfPosGetter.get_content_pos(pdf_path)
    print(pos)
    

    water = Watermark()
    # water.make_score_watermark(pdf_path,'张三',90)
    water.make_score_watermark_by_postions(pdf_path,'张三',90,pos)
